'''
1 读取集合 
2 读取实况
3 筛选快增算法
4 画图

# 这个脚本没用，是理解错需求写出来的脚本。
'''

import pandas as pd
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime,timedelta
import tqdm



RIstd = 7
tyname = f"dusurui RI-{RIstd}"
# members path
directory =  'D:\\met_data\\dusurui\\'    # 要筛选的文件开头
dates_name = os.listdir(directory)
dates_paths = [directory + date for date in dates_name]
obs_path = r'D:\met_data\ty_obs\dusurui_CMAobs.txt'


class tydat:
    def __init__(self,path,skiprows=2,extent=[100,170,0,60]):  # default - dusurui
        self.df=pd.read_csv(path,skiprows=skiprows,engine='python',sep=' & |\s',header=None) 
        self.colors=['#000000','#00ffff','#0000ff', '#FF8C00','#FF0000','#FF00FF']
        self.lon= self.df[8].values
        self.lat=self.df[9].values
        self.tp = self.df[14].values
        self.umax = self.df[13].values
        self.pmin = self.df[10].values
        self.extent=extent
        # 这里需要转成datetime对象
        self.time = np.array([datetime.strptime(str(t),'%Y%m%d%H') for t in self.df[0] ])
        
class tydat_CMA:
    def __init__(self,path,skiprows=2,extent=[100,170,0,60]):    
        dfs = pd.read_csv(path,engine='python',sep='\s+',header=None,dtype=str)
        tindex = []
        time = []
        for i in range(len(dfs[0])):    
            t = dfs[0][i] + dfs[1][i] + dfs[2][i] + dfs[3][i]
            t = datetime.strptime(t,"%Y%m%d%H")
            time.append(t)
            if np.isin(t.hour,[0,6,12,18]).item():
                tindex.append(i)
        df = pd.read_csv(path,engine='python',sep='\s+',header=None)
        # 2 是标准时刻，如0，6，12，18
        self.lon2 = np.array([df[5][i] for i in tindex]) 
        self.lat2 = np.array([df[6][i] for i in tindex])
        self.umax2 = np.array([df[8][i] for i in tindex])
        self.pmin2 = np.array([df[7][i] for i in tindex])
        self.time2 = np.array([time[i] for i in tindex])
        # all time
        self.lon = np.array(df[5])
        self.lat = np.array(df[6])
        self.umax = np.array(df[8])
        self.pmin = np.array(df[7])
        self.time = np.array(time)
        self.extent=extent


class tydat_JMA:
    def __init__(self,path,skiprows=1,extent=[100,170,0,60]):  # default - dusurui
        self.df=pd.read_csv(path,skiprows=skiprows,engine='python',header=None,sep="\s+") 
        self.colors=['#000000','#00ffff','#0000ff', '#FF8C00','#FF0000','#FF00FF']
        lat= self.df[3].values*0.1
        lon=self.df[4].values*0.1
        self.lon,self.lat=self.cntrl_lonlat(lon,lat)
        self.extent=extent
        # 这里需要转成datetime对象
        self.time = [datetime.strptime("20"+str(t),'%Y%m%d%H') for t in self.df[0] ]
        self.pmin = self.df[5]
        self.umax = self.df[6]*0.514  #kt to m/s
        
    def cntrl_lonlat(self,lon,lat):
        lon = lon[(lon>=0) & (lon <=360)] #去异常   有e25次方
        lon=np.where(lon>180,180,lon)  #防止左右边界相连
        return lon,lat
    



def count_rapidgrow(umax,t):
    '''
    umax:每个时次最大风速
    t:以datetime对象表示的时次
    return 一个true false对象
    '''
    i=0
    record = np.zeros(umax.shape,dtype=bool)
    if RIstd == 15:
        while t[i]+timedelta(days=1)<=t[-1]:
            trange = (t>=t[i])&(t<=t[i]+timedelta(days=1))
            urange = umax[trange]
            gap = max(urange)-urange[0]
            if gap>=RIstd:
                record[i]=1
            else:
                record[i]=0
            i+=1
    elif RIstd==7:
        while t[i]+timedelta(hours=12)<=t[-1]:
            trange = (t>=t[i])&(t<=t[i]+timedelta(hours=12))
            urange = umax[trange]
            gap = max(urange)-urange[0]
            if gap>=RIstd:
                record[i]=1
            else:
                record[i]=0
            i+=1
    return record


# 有底图那种
def track(ty,ax,lon,lat,title=None,color='black',label=None):
    ax.add_feature(cfeature.COASTLINE,linewidth=0.5) #海岸线
    ax.add_feature(cfeature.BORDERS, linestyle=':')   #国界
    china_provinces = cfeature.NaturalEarthFeature(category='cultural',name='admin_1_states_provinces_lines', scale='10m', facecolor='none')
    ax.add_feature(china_provinces, edgecolor='black', linewidth=0.5) #省界
    ax.set(xlim=(ty.extent[0],ty.extent[1]))
    ax.set(ylim=(ty.extent[2],ty.extent[3]))    
    #这种设定方法适用于这样子给出范围。
    ax.set_xticks(np.linspace(ty.extent[0], ty.extent[1]+5, 5))
    ax.set_yticks(np.linspace(ty.extent[2], ty.extent[3]+5, 5))
    ax.xaxis.set_major_formatter(LongitudeFormatter(number_format='.1f'))
    ax.xaxis.set_minor_locator(plt.MultipleLocator(1))
    ax.yaxis.set_major_formatter(LatitudeFormatter(number_format='.1f'))
    ax.yaxis.set_minor_locator(plt.MultipleLocator(1))
    ax.tick_params(axis='both', labelsize=6, direction='out')
    ###
    ax.plot(lon,lat,marker="o",transform=ccrs.PlateCarree(),color=color,label=label,markersize=6)





# 起报时间循环
dt_str=""
for i in tqdm.tqdm(range(len(dates_paths)), desc="Processing dates"):
    members_paths = [os.path.join(dates_paths[i], f) for f in os.listdir(dates_paths[i]) if f != "TRACK_ID_0" and f.startswith("TRACK_ID")]   # 打印所有符合后缀的文件名
    era_path = [os.path.join(dates_paths[i], f) for f in os.listdir(dates_paths[i]) if f.startswith("TRACK_ID_0")][0]  # 打印所有符合后缀的文件名
    fig, ax = plt.subplots(1, 1, figsize=(10, 10), subplot_kw={'projection': ccrs.PlateCarree()})
    # members
    colors = plt.cm.rainbow(np.linspace(0, 1, len(members_paths))) 
    labels = [f"ID {i}" for i in np.arange(len(members_paths))]
    for j in range(len(members_paths)):
        member = tydat(members_paths[j])
        record = count_rapidgrow(member.umax, member.time)
        if np.sum(record) != 0.0:
            lon = member.lon[record == 1]
            lat = member.lat[record == 1]
            a = np.where(record==1)[0][0]
            dt = member.time[a]-member.time[0]
            # dt_str = f" {dt.days}d{int(dt.seconds/3600)}hours"
            dt_str = member.time[record==1][-1].strftime("%d%H")
            track(member, ax, lon, lat,color=colors[j],label=labels[j]+"-"+dt_str)
    
    
    ### era5
    era = tydat(era_path)
    record = count_rapidgrow(era.umax, era.time)
    if np.sum(record) != 0.0:
        lon = era.lon[record == 1]
        lat = era.lat[record == 1]
        a = np.where(record==1)[0][0]
        dt = era.time[a]-member.time[0]
        # dt_str = f"{dt.days}d {int(dt.seconds/3600)}hours"
        dt_str =  era.time[record==1][-1].strftime("%d%H")
        track(era, ax, lon, lat,color="#4B0082",label="era-"+dt_str)
        
    obs = tydat_CMA(obs_path)
    record = count_rapidgrow(obs.umax, obs.time)
    if np.sum(record) != 0.0:
        lon = obs.lon[record == 1] 
        lat = obs.lat[record == 1]
        a = np.where(record==1)[0][0]
        dt = obs.time[a]-member.time[0]
        # dt_str = f"{dt.days}d {int(dt.seconds/3600)}hours"
        dt_str = obs.time[record==1][-1].strftime("%d%H")
        track(obs, ax, lon, lat,color="#2ca02c",label="obs-"+dt_str)
    
    # obs = tydat_JMA(obs_path)
    # record = count_rapidgrow(obs.umax, obs.time)
    # if np.sum(record) != 0.0:
    #     lon = obs.lon[record == 1]
    #     lat = obs.lat[record == 1]
    #     track(obs, ax, lon, lat,color="blue",label="obs")

    ax.set_title(dates_name[i]+f"{tyname}")
    ax.legend()
    # plt.savefig(f"D:\\met_data\\pics\\{dates_name[i]}.png") ; plt.close()
    plt.show()